#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2025-07-14

"""
【节点名称】：
    SMS2DTargetsToROS2
【订阅类型】：
    spirecv_msgs::2DTargets
【发布类型】：
    None
"""
import os.path
import threading
import time
import cv2
import argparse
from queue import Queue
from typing import Union
from spirems import Publisher, Subscriber, cvimg2sms, sms2cvimg, BaseNode, get_extra_args
from sv2_msgs.msg import Target, Targets, ROI
from rclpy.node import Node
import rclpy


class SMS2DTargetsToROS2(threading.Thread, BaseNode, Node):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown: bool = True,
        **kwargs
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=sms_shutdown,
            **kwargs
        )
        Node.__init__(self, self.__class__.__name__)

        self.target_topic = self.get_param("target_topic", "")
        self.params_help()

        input_topic = self.target_topic if len(self.target_topic) > 0 else '/' + job_name + '/detector/results'

        self._results_reader = Subscriber(
            input_topic, 'spirecv_msgs::2DTargets', self.results_callback,
            ip=ip, port=port
        )
        self._ros_pub = self.create_publisher(Targets, input_topic, 10)

        self.job_queue = Queue()
        self.queue_pool.append(self.job_queue)

        self.start()

    def release(self):
        BaseNode.release(self)

    def results_callback(self, msg):
        self.job_queue.put(msg)

    def run(self):
        while self.is_running():
            msg = self.job_queue.get(block=True)
            if msg is None:
                break

            # print(msg)
            tgts = Targets()
            tgts.header.stamp.sec = int(msg['timestamp'])
            tgts.header.stamp.nanosec = int((msg['timestamp'] - tgts.header.stamp.sec) * 1e9)
            tgts.frame_id = int(msg['img_id'])
            tgts.height, tgts.width = int(msg['height']), int(msg['width'])
            tgts.fps = float(msg['fps'])
            tgts.fov_x, tgts.fov_y = float(msg['fov_x']), float(msg['fov_y'])
            for r in msg['rois']:
                roi = ROI()
                roi.x1 = int(r[0])
                roi.y1 = int(r[1])
                roi.x2 = int(r[0] + r[2])
                roi.y2 = int(r[1] + r[3])
            for t in msg['targets']:
                tgt = Target()
                if 'bbox' in t:
                    tgt.cx = t['bbox'][0] + t['bbox'][2] / 2.0
                    tgt.cy = t['bbox'][1] + t['bbox'][3] / 2.0
                    tgt.w = float(t['bbox'][2])
                    tgt.h = float(t['bbox'][3])
                    tgt.los_ax = ((tgt.cx / tgts.width) - 0.5) * msg['fov_x']
                    tgt.los_ay = ((tgt.cy / tgts.height) - 0.5) * msg['fov_y']
                if 'score' in t:
                    tgt.score = float(t['score'])
                else:
                    tgt.score = 1.0
                if 'category_name' in t:
                    tgt.category = t['category_name']
                if 'category_id' in t:
                    tgt.category_id = int(t['category_id'])
                if 'pos' in t:
                    tgt.px = float(t['pos'][0])
                    tgt.py = float(t['pos'][1])
                    tgt.pz = float(t['pos'][2])
                if 'yaw_a' in t:
                    tgt.yaw_a = float(t['yaw_a'])
                tgts.targets.append(tgt)

            # print(tgts)
            self._ros_pub.publish(tgts)

        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        type=str,
        default='default_params.json',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    args, unknown_args = parser.parse_known_args()
    if not os.path.isabs(args.config):
        current_path = os.path.abspath(__file__)
        params_dir = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], 'params', 'spirecv2')
        args.config = os.path.join(params_dir, args.config)
    print("--config:", args.config)
    print("--job-name:", args.job_name)
    extra = get_extra_args(unknown_args)
    rclpy.init(args=None)

    node = SMS2DTargetsToROS2(args.job_name, param_dict_or_file=args.config, ip=args.ip, port=args.port, **extra)
    node.spin()
